from pymongo import MongoClient
from pymongo.errors import BulkWriteError
from conf.settings import CLOUDMONITOR_MONGODB, CLOUDMONITOR_DB
import pymongo
import logging
from enum import Enum

logger = logging.getLogger(__name__)


class ScanProfile(Enum):
    """enum ScanProfile {
        SP_DEFAULT = 0;
        SP_FULL = 1;
        SP_HIGH_RISK = 2;
        SP_SQL_INJECTION = 3;
        SP_XSS = 4;
        SP_NORMAL = 5;
        SP_QUICK = 6;
        }

    Args:
        Enum ([type]): [description]
    """
    high = 2                     # 高危漏洞检测
    medium = 1                   # 全面深入检测
    low = 1                      # 全面深入检测
    file_upload = 14
    directory_listing = 11
    xss = 4
    directory_traversal = 12
    code_execution = 7
    csrf = 10
    file_inclusion = 13
    injection = 16
    configuration = 9
    information_disclosure = 15
    command_execution = 8
    malware = 17
    webvul = 18


def update_vulsettings(col):
    query = {
        "taskGroupType": "web_security"
    }
    jobs_iter = col.find(query)
    bulk_list = []
    modified_count = 0
    for data in jobs_iter:
        collect = data.get("vulSettings", {}).get("collect", {})
        vul = collect.get("vul", [])
        _id = data.get("_id")
        new_vul = []
        vulType = data.get("vulSettings", {}).get("vulType", "")
        for name in vul:
            if hasattr(ScanProfile, name):
                new_vul.append(str(getattr(ScanProfile, name).value))
            else:
                if not vulType and not name.isdigit():
                    vulType = "plugIn"
                break
        update_data = {
            "vulSettings.vulType": vulType
        }
        if new_vul:
            collect["vul"] = new_vul
            update_data["vulSettings.collect.vul"] = new_vul
        bulk_list.append(pymongo.UpdateOne({"_id":_id}, {"$set": update_data}, upsert=False))
        if len(bulk_list) == 300:
            try:
                result = col.bulk_write(bulk_list, ordered=False)
            except BulkWriteError as bwe:
                logger.exception(bwe.details)
            else:
                modified_count += result.modified_count
                logger.info(f"{modified_count} modified")
                bulk_list = []
    if bulk_list:
        try:
            result = col.bulk_write(bulk_list, ordered=False)
        except BulkWriteError as bwe:
            logger.exception(bwe.details)
        else:
            modified_count += result.modified_count
            logger.info(f"{modified_count} modified")


def main():
    logging.basicConfig(level=logging.DEBUG)
    assert CLOUDMONITOR_MONGODB and CLOUDMONITOR_DB
    db = MongoClient(CLOUDMONITOR_MONGODB)[CLOUDMONITOR_DB]
    col = db["asset_tasks"]
    update_vulsettings(col)

if __name__ == "__main__":
    main()